import json
import re

import cv2
import joblib
import numpy as np
from flask import Flask, request, jsonify
from skimage.feature import hog


# 图像预处理函数：调整图像大小，灰度化，归一化
def preprocess_image(image, target_size=(128, 128)):
    image_resized = cv2.resize(image, target_size)
    gray_image = cv2.cvtColor(image_resized, cv2.COLOR_BGR2GRAY)
    normalized_image = gray_image / 255.0
    return normalized_image


# 提取图像的 HOG 特征
def extract_hog_features(images):
    hog_features = []
    for image in images:
        hog_feature = hog(image, orientations=9, pixels_per_cell=(8, 8), cells_per_block=(2, 2), transform_sqrt=True)
        hog_features.append(hog_feature)
    return np.array(hog_features)


# 加载模型
def load_model(model_path):
    return joblib.load(model_path)


# 加载已训练好的模型和标签映射
def load_model_and_labels(model_path, label_map):
    svm_model = load_model(model_path)
    return svm_model, label_map


# 植物图像识别
def predict_plant(image, svm_model, label_map):
    preprocessed_image = preprocess_image(image)
    hog_feature = extract_hog_features([preprocessed_image])
    prediction = svm_model.predict(hog_feature)
    plant_label = label_map[prediction[0]]
    return plant_label


# 加载mushroom.json
def load_mushroom_data():
    try:
        with open('mushroom.json', 'r', encoding='utf-8') as f:
            return json.load(f)
    except Exception as e:
        print(f"Error loading mushroom data: {e}")
        return []


# mushroom.json的列表中找到type_name对应一项
def find_mushroom_by_type(type_name, mushroom_data):
    return next(
        (item for item in mushroom_data if item["type_name"].lower() == type_name.lower()),
        None
    )


def create_flask_app():
    app = Flask(__name__)

    # 加载模型和标签映射
    model_path = "model/svm_model_1221.joblib"  # 模型路径
    # 标签映射
    label_map = {0: 'Beech Mushroom', 1: 'Black Trumpet Mushroom', 2: 'Button Mushroom', 3: 'Chanterelle Mushroom',
                 4: 'Crimini Mushroom', 5: 'Enoki Mushroom', 6: 'Hedgehog Mushroom', 7: 'King Oyster Mushroom',
                 8: 'Maitake Mushroom', 9: 'Morel Mushroom', 10: 'Oyster Mushroom', 11: 'Porcini Mushroom',
                 12: 'Portobello Mushroom', 13: 'Shiitake Mushroom'}
    svm_model, label_map = load_model_and_labels(model_path, label_map)

    # 加载蘑菇数据
    mushroom_data = load_mushroom_data()

    @app.route('/predict', methods=['POST'])
    def predict():
        try:
            if 'image' not in request.files:
                return jsonify({'status': 'error', 'message': 'No image file provided'}), 400

            file = request.files['image']
            if file.filename == '':
                return jsonify({'status': 'error', 'message': 'No selected file'}), 400

            # 读取上传的图像
            image = cv2.imdecode(np.fromstring(file.read(), np.uint8), cv2.IMREAD_COLOR)
            if image is None:
                return jsonify({'status': 'error', 'message': 'Invalid image'}), 400

            # 获取预测结果
            predicted_type = predict_plant(image, svm_model, label_map)
            # 查找对应的蘑菇详细信息
            mushroom_info = find_mushroom_by_type(predicted_type, mushroom_data)
            if mushroom_info:
                return jsonify({
                    'status': 'success',
                    'predicted_type': predicted_type,
                    'mushroom_details': mushroom_info
                })
            else:
                return jsonify({
                    'status': 'error',
                    'message': f'No detailed information found for mushroom type: {predicted_type}'
                }), 404

        except Exception as e:
            return jsonify({
                'status': 'error',
                'message': f'Error processing request: {str(e)}'
            }), 500

    @app.route('/search', methods=['GET'])
    def search():
        print(mushroom_data)
        keyword = request.args.get('keyword', '')
        # 如果 keyword 为空，返回所有数据
        if not keyword:
            return jsonify({
                'status': 'success',
                'data': mushroom_data,
            })

        # 模糊匹配关键词
        results = [mushroom for mushroom in mushroom_data if
                   re.search(re.escape(keyword), mushroom['type_chinese_name'], re.IGNORECASE)]

        return jsonify({
            'status': 'success',
            'data': results,
        })

    return app


if __name__ == '__main__':
    app = create_flask_app()
    app.run(debug=True, host='0.0.0.0', port=5000)
